import os
import sys
import struct


# def detect_dos_sign(path):
#     try:
#         fileHandle = open(path, 'rb')
#         data = struct.unpack('h', fileHandle.read(2))[0]
#         fileHandle.close()
#         if data == 0x5A4D:
#             return True
#     except:
#         return False
#
#
# def detect_nt_sign(path):
#     try:
#         fileHandle = open(path, 'rb')
#         fileHandle.seek(0x3C, 0)
#         data1 = struct.unpack('h', fileHandle.read(2))[0]
#         # fileHandle.close()
#         # fileHandle = open(path, 'rb')
#         fileHandle.seek(data1, 0)
#         data2 = struct.unpack('h', fileHandle.read(2))[0]
#         fileHandle.close()
#         if data2 == 0x4550:
#             return True
#     except:
#         return False


def detect_pe(path):
    try:
        fileHandle = open(path, 'rb')
        data = struct.unpack('h', fileHandle.read(2))[0]
        if data != 0x5A4D:
            fileHandle.close()
            return False
        fileHandle.seek(0x3C, 0)
        data1 = struct.unpack('h', fileHandle.read(2))[0]
        fileHandle.seek(data1, 0)
        data2 = struct.unpack('h', fileHandle.read(2))[0]
        fileHandle.close()
        if data2 == 0x4550:
            return True
        else:
            return False
    except:
        return False


def detect_elf_sign(path):
    try:
        fileHandle = open(path, 'rb')
        data_id = struct.unpack('i', fileHandle.read(4))[0]
        fileHandle.close()
        if data_id == 0x464C457F:
            return True
    except Exception as e:
        return False


# def find_all_executable_files(filepath):
#     pe_exetensions = ['.exe', '.dll', '.ocx', '.sys', '.com', '.lib']
#     elf_extensions = ['.o', '.out', '.so']
#     for root, dirs, files in os.walk(filepath):
#         for file in files:
#             for pe_extension in pe_exetensions:
#                 if file.endswith(pe_extension):
#                     fullname = os.path.join(root, file)
#                     yield fullname
#             for elf_extension in elf_extensions:
#                 if file.endswith(elf_extension):
#                     fullname = os.path.join(root, file)
#                     yield fullname


def find_all_executable_files(filepath):
    for root, dirs, files in os.walk(filepath):
        for file in files:
            fullname = os.path.join(root, file)
            # if detect_dos_sign(fullname) and detect_nt_sign(fullname):
            if detect_pe(fullname):
                yield fullname
            if detect_elf_sign(fullname):
                yield fullname


def get_executable_files(package):
    executable_file_list = []
    for executable_file in find_all_executable_files(package):
        # print(python_file)
        executable_file_list.append(executable_file)
    return executable_file_list


# pe1 = detect_dos_sign('test/filetype_test/PEtest.exe')
# pe2 = detect_nt_sign('test/filetype_test/PEtest.exe')
# elf=detect_elf_sign('test/filetype_test/Question1')
# print(pe1, pe2, elf)

package = sys.argv[1]
executable_file_list = get_executable_files(package)
if executable_file_list:
    if len(executable_file_list) == 1:
        print('Package %s has 1 executable file!!!' % package)
        print('The executable file is %s!' % executable_file_list[0])
    else:
        print('Package %s has %s executable files!!!' % (package, str(len(executable_file_list))))
        print('There are the following files:')
        for executable_file in executable_file_list:
            print(executable_file)
